# 数据库操作
from pymongo import MongoClient, UpdateOne

"""
数据结构：实际使用时需考虑任务是否有唯一性要求、_id设计以及其他必须字段。
    {
        "_id": "test_func_6_test_func",
        "interval": NumberInt("6"),
        "plugin_name": "test_func",
        "status": "dormant",
        "task_id": "test_func"
    }
测试数据如上，实际项目需确认task的唯一标识
"""


class Mongo:
    # 单例模式
    _instance = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self):
        if not hasattr(self, "client"):
            self.client = MongoClient('mongodb://localhost:27017/')
            self.db = self.client['test_automation']
            self.collection = self.db["tasks"]

    def get_client(self):
        """
        提供外部访问collection
        :return:
        """
        return self.collection

    def close(self):
        self.client.close()

    def exits_update_and_insert(self, data, bulk=False, ordered=False):
        """
        存在则更新，不存在则插入
        :param data: 数据
        :param bulk: 是否使用批量插入
        :param ordered: 有序执行，一条报错，后面不再执行
                        无序执行，一条报错，其他不受影响
        :return:

        """
        if bulk:
            bulk_docs = []
            for rd in data:
                bulk_docs.append(
                    UpdateOne(
                        {"_id": rd.get("_id") or f"{'_'.join(map(str, list(rd.values())[:-1]))}"},
                        {"$set": rd},
                        upsert=True
                    )
                )
            result = self.collection.bulk_write(bulk_docs, ordered=ordered).bulk_api_result
            print(f"exits_update_and_insert result: {result}")
        else:
            for rd in data:
                r = self.collection.update_one(
                    {"_id": rd.get("_id") or f"{'_'.join(map(str, list(rd.values())[:-1]))}"},
                    {"$set": rd},
                    upsert=True
                )
                print(f"exits_update_and_insert: 操作完成。 "
                      f"task id: {rd['task_id']}, "
                      f"modified_count: {r.modified_count}")

    def exits_do_nothing_and_insert(self, data, bulk=False, ordered=False):
        """
        存在则不做任何操作
        不存在则插入
        :param data: 数据
        :param bulk: 是否批量
        :param ordered: 是否有序
        :return:
        """
        if bulk:
            bulk_docs = []
            for rd in data:
                bulk_docs.append(
                    UpdateOne(
                        {"_id": rd.get("_id") or f"{'_'.join(map(str, list(rd.values())[:-1]))}"},
                        {"$setOneInsert": rd},
                        upsert=True
                    )
                )
                result = self.collection.bulk_write(bulk_docs, ordered=ordered).bulk_api_result
                print(result)
        else:
            for rd in data:
                self.collection.update_one(
                    {"_id": rd.get("_id") or f"{'_'.join(map(str, list(rd.values())[:-1]))}"},
                    {"$setOneInsert": rd},
                    upsert=True
                )

    def find(self, query):
        """
        查询
        :param query:
        :return:
        """
        r = list(self.collection.find(query))
        return r

    def remove(self, query):
        """
        删除
        :param query:
        :return:
        """
        r = self.collection.delete_one(query)
        return r


# 全局实例
MGClient = Mongo()
